import threading
import queue
import time
import logging
from concurrent.futures import ThreadPoolExecutor, as_completed
from tkinter import ttk, messagebox, scrolledtext, filedialog
from typing import List, Dict, Optional, Callable
import pandas as pd
from dataclasses import dataclass
from enum import Enum

# 修复导入问题
try:
    from crawler import DouyinCrawler
    from logger import logger
except ImportError:
    # 如果相对导入失败，尝试绝对导入
    import sys
    import os
    sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    try:
        from crawler import DouyinCrawler
        from logger import logger
    except ImportError:
        # 创建一个简单的logger作为后备
        import logging
        logger = logging.getLogger('threaded_crawler')
        logger.setLevel(logging.INFO)
        console_handler = logging.StreamHandler()
        console_handler.setLevel(logging.INFO)
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        console_handler.setFormatter(formatter)
        logger.addHandler(console_handler)


class TaskStatus(Enum):
    """任务状态枚举"""
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


@dataclass
class CrawlTask:
    """爬取任务数据类"""
    task_id: str
    url: str
    scroll_times: int
    status: TaskStatus = TaskStatus.PENDING
    result: Optional[List[Dict]] = None
    error: Optional[str] = None
    start_time: Optional[float] = None
    end_time: Optional[float] = None


class ThreadSafeDataManager:
    """线程安全的数据管理器"""
    
    def __init__(self):
        self._data_lock = threading.Lock()
        self._comments_data = []
        self._task_results = {}
        self._progress_callbacks = []
        
    def add_comments(self, comments: List[Dict]):
        """线程安全地添加评论数据"""
        with self._data_lock:
            self._comments_data.extend(comments)
        self._notify_progress(comments)
    
    def get_all_comments(self) -> List[Dict]:
        """线程安全地获取所有评论数据"""
        with self._data_lock:
            return self._comments_data.copy()
    
    def clear_comments(self):
        """清空评论数据"""
        with self._data_lock:
            self._comments_data.clear()
    
    def add_progress_callback(self, callback: Callable):
        """添加进度回调函数"""
        self._progress_callbacks.append(callback)
    
    def _notify_progress(self, new_comments: List[Dict]):
        """通知进度更新"""
        for callback in self._progress_callbacks:
            try:
                callback(new_comments)
            except Exception as e:
                logger.error(f"进度回调函数执行失败: {e}")


class MultiThreadedDouyinCrawler:
    """多线程抖音评论爬虫"""
    
    def __init__(self, max_workers: int = 3):
        self.max_workers = max_workers
        self.data_manager = ThreadSafeDataManager()
        self.task_queue = queue.Queue()
        self.active_tasks = {}
        self._stop_event = threading.Event()
        self._executor = None
        self._main_thread = None
        self._progress_callback = None
        self._gui_root = None
        # 新增：用于存储所有活跃的爬虫实例，实现外部停止控制
        self._active_crawlers = []
        self._crawlers_lock = threading.Lock()
        self.extract_type = "pinglun"
        
    def set_gui_callback(self, callback: Callable, gui_root=None):
        """设置GUI回调函数和根窗口"""
        self._progress_callback = callback
        self._gui_root = gui_root
        # 直接将回调传递给数据管理器
        self.data_manager.add_progress_callback(self._handle_progress_update)
    
    def _handle_progress_update(self, new_comments: List[Dict]):
        """处理来自数据管理器的进度更新"""
        if self._progress_callback and self._gui_root:
            # 将更新请求放入GUI的队列中
            self._gui_root._update_queue.put(('new_comment', new_comments))

    def stop_all_tasks(self):
        """停止所有任务"""
        self._stop_event.set()
        logger.info("正在停止所有爬取任务...")
        
        # 获取所有活跃爬虫的副本以避免在迭代时被修改
        with self._crawlers_lock:
            crawlers_copy = self._active_crawlers.copy()
        
        # 通知所有正在进行的爬虫实例立即停止
        for crawler in crawlers_copy:
            try:
                crawler.stop_scroll = True
                logger.info("已向爬虫实例发送停止信号")
            except Exception as e:
                logger.error(f"向爬虫实例发送停止信号时出错: {e}")
    
    def close(self):
        """关闭爬虫，释放资源"""
        logger.info("正在关闭多线程爬虫...")
        
        # 停止所有任务
        self.stop_all_tasks()
        
        # 等待主线程结束
        if self._main_thread and self._main_thread.is_alive():
            try:
                self._main_thread.join(timeout=5)  # 等待最多5秒
            except Exception as e:
                logger.error(f"等待主线程结束时出错: {e}")
        
        # 关闭线程池
        if self._executor:
            try:
                self._executor.shutdown(wait=False)  # 立即关闭，不等待正在执行的任务
            except Exception as e:
                logger.error(f"关闭线程池时出错: {e}")
        
        logger.info("多线程爬虫已关闭")
    
    def clear_data(self):
        """清空所有数据"""
        self.data_manager.clear_comments()
        self.active_tasks.clear()
         # crawler也需要清理
        for crawler in self._active_crawlers:
            crawler.clear_data()

        logger.info("已清空所有数据")
    
    def add_crawl_task(self, url: str, scroll_times: int = 5) -> str:
        """添加爬取任务"""
        import uuid
        task_id = str(uuid.uuid4())
        task = CrawlTask(
            task_id=task_id,
            url=url,
            scroll_times=scroll_times
        )
        self.task_queue.put(task)
        self.active_tasks[task_id] = task
        logger.info(f"添加新任务: {task_id}")
        return task_id
    
    def get_task_status(self, task_id: str) -> Optional[CrawlTask]:
        """获取任务状态"""
        return self.active_tasks.get(task_id)
    
    def get_all_tasks(self) -> List[CrawlTask]:
        """获取所有任务"""
        return list(self.active_tasks.values())
    
    def get_progress_stats(self) -> Dict:
        """获取进度统计信息"""
        tasks = self.get_all_tasks()
        total = len(tasks)
        completed = sum(1 for task in tasks if task.status == TaskStatus.COMPLETED)
        running = sum(1 for task in tasks if task.status == TaskStatus.RUNNING)
        failed = sum(1 for task in tasks if task.status == TaskStatus.FAILED)
        
        return {
            'total': total,
            'completed': completed,
            'running': running,
            'failed': failed,
            'total_comments': len(self.data_manager.get_all_comments())
        }
    
    def _execute_crawl_task(self, task: CrawlTask) -> List[Dict]:
        """执行单个爬取任务"""
        task.status = TaskStatus.RUNNING
        task.start_time = time.time()
        
        crawler = None
        try:
            logger.info(f"开始执行任务 {task.task_id}: {task.url}")
            
            # 创建爬虫实例
            crawler = DouyinCrawler()
            
            # 如果GUI根窗口存在，设置浏览器关闭回调
            if self._gui_root and hasattr(self._gui_root, 'on_browser_closed'):
                crawler.set_browser_close_callback(self._gui_root.on_browser_closed)
            
            # 初始化浏览器
            if not crawler.init_browser():
                raise Exception("浏览器初始化失败")
            
            # 设置GUI回调函数
            if self._progress_callback:
                crawler.set_gui_callback(self._notify_single_comment)
            
            # 将爬虫实例添加到活跃列表中，以便可以被停止
            with self._crawlers_lock:
                self._active_crawlers.append(crawler)
            
            # 设置停止标志检查
            if self._stop_event.is_set():
                crawler.stop_scroll = True
            # 访问用户主页
            logger.info(f"打开用户主页: {task.url}")
            crawler.page.get(task.url)
            """显示登录提示对话框"""
            result = messagebox.askyesno("请登录", "浏览器已打开，界面验证成功后点击'是'开始采集用户，点击'否'取消操作。")
            if result:
                # 执行爬取
                if self.extract_type == "pinglun":
                    success = crawler.get_comments(task.url, task.scroll_times)
                elif self.extract_type == "guanzhu":
                    # 用户确认已登录，开始采集
                    success = crawler.get_following_comments(task.url, task.scroll_times)
            else:
                return []
           
                
            if success:
                comments = crawler.comments_data
                logger.info(f"任务 {task.task_id} 完成，获取到 {len(comments)} 条评论")
                
                task.status = TaskStatus.COMPLETED
                task.result = comments
                task.end_time = time.time()
                
                return comments
            else:
                raise Exception("爬取失败")
                
        except Exception as e:
            error_msg = str(e)
            logger.error(f"任务 {task.task_id} 失败: {error_msg}")
            
            task.status = TaskStatus.FAILED
            task.error = error_msg
            task.end_time = time.time()
            
            return []
        
        finally:
            # 从活跃列表中移除爬虫实例
            if crawler is not None:
                with self._crawlers_lock:
                    if crawler in self._active_crawlers:
                        self._active_crawlers.remove(crawler)
                
                # 确保浏览器被关闭
                try:
                    crawler.close()
                except:
                    pass
    
    def _notify_single_comment(self, comment_data: List[Dict]):
        """线程安全地处理单个评论更新"""
        self.data_manager.add_comments(comment_data)
    
    def start_crawling(self, blocking: bool = False):
        """开始爬取任务"""
        logger.info(f"开始多线程爬取，最大工作线程数: {self.max_workers}")
        
        if blocking:
            self._run_crawling_loop()
        else:
            self._main_thread = threading.Thread(target=self._run_crawling_loop)
            self._main_thread.daemon = True
            self._main_thread.start()
    
    def _run_crawling_loop(self):
        """运行爬取循环"""
        self._executor = ThreadPoolExecutor(max_workers=self.max_workers)
        
        try:
            futures = []
            
            # 提交所有任务
            while not self._stop_event.is_set():
                try:
                    # 从队列获取任务，超时1秒
                    task = self.task_queue.get(timeout=1)
                    
                    # 提交任务到线程池
                    future = self._executor.submit(self._execute_crawl_task, task)
                    futures.append((future, task))
                    
                except queue.Empty:
                    # 队列为空，检查是否还有正在执行的任务
                    if not futures:
                        break
                    continue
                except Exception as e:
                    logger.error(f"提交任务时出错: {e}")
                    continue
            
            # 如果收到停止信号，不再等待任务完成
            if self._stop_event.is_set():
                logger.info("收到停止信号，取消提交新任务并跳过等待剩余任务")
            else:
                # 否则正常等待所有任务完成
                logger.info(f"等待 {len(futures)} 个任务完成...")
                for future, task in futures:
                    try:
                        # 等待单个任务完成，设置超时
                        result = future.result(timeout=300)  # 5分钟超时
                        logger.info(f"任务 {task.task_id} 完成，获取到 {len(result)} 条评论")
                    except Exception as e:
                        logger.error(f"任务 {task.task_id} 执行超时或失败: {e}")
                        task.status = TaskStatus.FAILED
                        task.error = str(e)
                    
        except Exception as e:
            logger.error(f"爬取循环出错: {e}")
            
        finally:
            # 清理资源
            if self._executor:
                self._executor.shutdown(wait=True)
            
            logger.info("多线程爬取完成")
            
            # 通知GUI爬取完成
            if self._progress_callback and self._gui_root:
                try:
                    final_data = self.data_manager.get_all_comments()
                    self._gui_root.after(0, lambda: self._progress_callback(final_data, completed=True))
                except Exception as e:
                    logger.error(f"通知GUI完成时出错: {e}")
    
    def wait_for_completion(self, timeout: Optional[float] = None):
        """等待所有任务完成"""
        if self._main_thread and self._main_thread.is_alive():
            self._main_thread.join(timeout=timeout)
    
    def save_results(self, filename: str = None) -> str:
        """保存结果到文件"""
        if not filename:
            import datetime
            filename = f"comments_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
        
        comments = self.data_manager.get_all_comments()
        
        if not comments:
            logger.warning("没有数据可保存")
            return ""
        
        try:
            df = pd.DataFrame(comments)
            df.to_csv(filename, index=False, encoding='utf-8-sig')
            logger.info(f"数据已保存到 {filename}")
            return filename
        except Exception as e:
            logger.error(f"保存文件失败: {e}")
            return ""


# 全局爬虫实例
global_crawler = None


def get_global_crawler(max_workers: int = 3) -> MultiThreadedDouyinCrawler:
    """获取全局爬虫实例"""
    global global_crawler
    if global_crawler is None:
        global_crawler = MultiThreadedDouyinCrawler(max_workers)
    return global_crawler


def crawl_with_threads(video_urls: List[str], scroll_times: int = 5, max_workers: int = 3) -> List[Dict]:
    """使用多线程爬取多个视频的评论"""
    crawler = get_global_crawler(max_workers)
    
    # 清空之前的数据
    crawler.clear_data()
    
    # 添加所有任务
    for url in video_urls:
        crawler.add_crawl_task(url, scroll_times)
    
    # 开始爬取（阻塞模式）
    crawler.start_crawling(blocking=True)
    
    # 返回所有数据
    return crawler.data_manager.get_all_comments()